import cv2
import numpy as np
import matplotlib.pyplot as plt
from camera import get_photo


def object_detection(camera_index=0):
    _, frame = get_photo(camera_index)

    # define range of blue color in HSV
    blue_lower = np.array([100, 100, 100])
    blue_upper = np.array([124, 255, 255])

    img_HSV = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    mask = cv2.inRange(img_HSV, lowerb=blue_lower, upperb=blue_upper)
    img_ball = cv2.bitwise_and(frame, frame, mask=mask)

    # 转换为二值图
    img_ball = cv2.cvtColor(img_ball, cv2.COLOR_BGR2GRAY)
    __, img_bin = cv2.threshold(img_ball, 10, 255, cv2.THRESH_BINARY)

    # 轮廓检测
    contours, __ = cv2.findContours(img_bin, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    # 寻找面积最大的轮廓，并获取最小外接矩形
    area = []
    for i in range(len(contours)):
        mianji = cv2.contourArea(contours[i])
        area.append(mianji)
    max_area_num = np.argmax(area)
    cnt = contours[max_area_num]
    print(cnt)
    x, y, w, h = cv2.boundingRect(cnt)

    # 绘制图像中心点到球中心的的连线
    x1 = int(x + w / 2)
    y1 = int(y + h / 2)
    # 调试用
    # Center_point_x = int(frame.shape[0] / 2)
    # Center_point_y = int(frame.shape[1] / 2)
    # img_juxing = cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)  # 画出最小外接矩形
    # img_juxing = cv2.line(img_juxing, (x1, y1), (Center_point_x, Center_point_y),
    #                       (0, 0, 255), 2)
    # plt.subplot(2,2,1)
    # plt.imshow(frame)
    # plt.subplot(2,2,2)
    # plt.imshow(mask)
    # plt.subplot(2,2,3)
    # plt.imshow(img_ball)
    # plt.subplot(2,2,4)
    # plt.imshow(img_juxing)
    # plt.show()
    return x1, y1


if __name__ == '__main__':
    center_x, center_y = object_detection(0)
    print(center_x, center_y)
